import tensorflow as tf
from transformers import GPT2Tokenizer, TFGPT2LMHeadModel
from flask import Flask, jsonify, request

model_dir = "C:/01.project/ipython/gpt-2/models/gpt2"

# 设置 TensorFlow 运行时使用的 GPU 设备
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        # 只使用第一块 GPU
        tf.config.experimental.set_visible_devices(gpus[0], 'GPU')
        tf.config.experimental.set_memory_growth(gpus[0], True)
    except RuntimeError as e:
        print(e)

# 加载 GPT-2 tokenizer
tokenizer = GPT2Tokenizer.from_pretrained(model_dir)

# 加载 GPT-2 模型
model = TFGPT2LMHeadModel.from_pretrained(model_dir)

# 将模型转换为 GPU 上运行
model = model.to('cuda')

# 创建 Flask 实例
app = Flask(__name__)

# 定义 API 接口
@app.route('/generate', methods=['POST'])
def generate_text():
    # 从请求中获取文本
    input_text = request.get_json()['text']

    # 使用 tokenizer 对文本进行编码
    input_ids = tokenizer.encode(input_text, return_tensors='tf')

    # 使用模型生成文本
    output_ids = model.generate(input_ids.to('cuda'), max_length=100, num_return_sequences=1)

    # 使用 tokenizer 对生成的文本进行解码
    output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)

    # 返回生成的文本
    return jsonify({'generated_text': output_text})

# 启动应用
if __name__ == '__main__':
    app.run(debug=True)
